#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
小市值轮动策略回测系统
基于北交所数据，选择总市值最小的N只股票，每M个交易日调仓
"""

import sqlite3
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from datetime import datetime, timedelta
import logging
import json
import os
from typing import List, Dict, Tuple, Optional
import warnings

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('small_cap_strategy.log', encoding='utf-8'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)


class SmallCapRotationStrategy:
    """小市值轮动策略类"""

    def __init__(self, db_path: str = "bse_data.db"):
        """
        初始化策略

        Args:
            db_path: 数据库路径
        """
        self.db_path = db_path
        self.conn = None

        # 策略参数
        self.N = 5  # 选择股票数量
        self.M = 10  # 调仓周期（交易日）
        self.initial_capital = 1000000  # 初始资金
        self.commission_rate = 0.002  # 交易手续费（千分之2）
        self.slippage_rate = 0.001  # 滑点（千分之1）

        # 过滤条件配置
        self.filter_conditions = {
            'pe_ttm': {'min': 0, 'max': None},        # PE-TTM > 0
            # 'pb': {'min': None, 'max': None},          # PB范围
            # 'ps_ttm': {'min': None, 'max': None},      # PS-TTM范围
             'dyr': {'min': 0.01, 'max': None},         # 股息率范围
            # 'to_r': {'min': None, 'max': None},        # 换手率范围
            # 'mc': {'min': None, 'max': None},          # 市值范围
        }

        # 上市天数过滤条件
        self.min_trading_days = 250  # 默认上市至少250个交易日

        # 北证50成分股过滤条件
        self.filter_bse50_constituents = False  # 是否只选择北证50成分股

        # 回测结果存储
        self.portfolio_value = []
        self.positions = {}  # 持仓信息
        self.trades = []  # 交易记录
        self.daily_returns = []

        # 数据缓存
        self.price_data = None
        self.equity_data = None
        self.fundamental_data = None
        self.merged_data = None

        self._connect_database()
        self._load_data()

    def _connect_database(self):
        """连接数据库"""
        try:
            self.conn = sqlite3.connect(self.db_path)
            logger.info(f"数据库连接成功: {self.db_path}")
        except Exception as e:
            logger.error(f"数据库连接失败: {e}")
            raise

    def _load_data(self):
        """加载股票价格、股本和基本面数据"""
        logger.info("开始加载数据...")

        # 加载K线数据
        price_query = """
        SELECT stock_code, date, close, open, high, low, volume
        FROM candlestick
        WHERE close IS NOT NULL
        ORDER BY stock_code, date
        """
        self.price_data = pd.read_sql_query(price_query, self.conn)

        # 加载股票信息（获取上市日期）
        stock_info_query = """
        SELECT stock_code, ipo_date
        FROM stock_info
        WHERE ipo_date IS NOT NULL
        """
        self.stock_info = pd.read_sql_query(stock_info_query, self.conn)

        # 加载股本数据
        equity_query = """
        SELECT stock_code, date, capitalization
        FROM equity_change
        WHERE capitalization IS NOT NULL
        ORDER BY stock_code, date
        """
        self.equity_data = pd.read_sql_query(equity_query, self.conn)

        # 加载基本面数据
        fundamental_query = """
        SELECT stock_code, date, pe_ttm, pb, ps_ttm, dyr, to_r, mc, cmc, ecmc
        FROM fundamental_data
        WHERE pe_ttm IS NOT NULL OR pb IS NOT NULL OR ps_ttm IS NOT NULL
        ORDER BY stock_code, date
        """
        self.fundamental_data = pd.read_sql_query(fundamental_query, self.conn)

        # 加载北证50成分股变更数据
        bse50_query = """
        SELECT date, constituents, change_type
        FROM bse50_constituents_changes
        ORDER BY date
        """
        self.bse50_constituents_data = pd.read_sql_query(bse50_query, self.conn)

        logger.info(f"加载价格数据: {len(self.price_data)} 条记录")
        logger.info(f"加载股本数据: {len(self.equity_data)} 条记录")
        logger.info(f"加载基本面数据: {len(self.fundamental_data)} 条记录")
        logger.info(f"加载北证50成分股数据: {len(self.bse50_constituents_data)} 条记录")

        # 检查数据质量
        self._check_data_quality()

        # 合并数据
        self._merge_data()

    def _check_data_quality(self):
        """检查数据质量"""
        # 检查价格数据
        null_prices = self.price_data['close'].isnull().sum()
        if null_prices > 0:
            logger.warning(f"发现 {null_prices} 条空价格数据")

        # 检查股本数据
        null_equity = self.equity_data['capitalization'].isnull().sum()
        if null_equity > 0:
            logger.warning(f"发现 {null_equity} 条空股本数据")

        # 检查基本面数据
        if self.fundamental_data is not None and len(self.fundamental_data) > 0:
            null_fundamental = self.fundamental_data.isnull().sum().sum()
            logger.info(f"基本面数据完整性检查:")
            for col in ['pe_ttm', 'pb', 'ps_ttm', 'dyr', 'to_r', 'mc']:
                if col in self.fundamental_data.columns:
                    count = self.fundamental_data[col].notna().sum()
                    logger.info(f"  {col}: {count} 条有效数据")
        else:
            logger.warning("没有基本面数据或基本面数据为空")

        # 检查股票数量
        unique_stocks_price = self.price_data['stock_code'].nunique()
        unique_stocks_equity = self.equity_data['stock_code'].nunique()
        unique_stocks_fundamental = self.fundamental_data['stock_code'].nunique() if self.fundamental_data is not None else 0

        logger.info(f"价格数据股票数: {unique_stocks_price}")
        logger.info(f"股本数据股票数: {unique_stocks_equity}")
        logger.info(f"基本面数据股票数: {unique_stocks_fundamental}")

        if unique_stocks_price < 50:
            logger.warning(f"价格数据股票数量较少: {unique_stocks_price}只")

        if unique_stocks_fundamental == 0:
            logger.warning("没有基本面数据，过滤条件将不会生效")

    def _merge_data(self):
        """合并价格、股本和基本面数据"""
        logger.info("合并价格、股本和基本面数据...")

        # 转换日期格式
        self.price_data['date'] = pd.to_datetime(self.price_data['date']).dt.tz_localize(None)
        self.equity_data['date'] = pd.to_datetime(self.equity_data['date']).dt.tz_localize(None)

        if self.fundamental_data is not None:
            self.fundamental_data['date'] = pd.to_datetime(self.fundamental_data['date']).dt.tz_localize(None)

        # 按股票代码分组，向前填充股本数据
        equity_filled = []
        for stock_code, group in self.equity_data.groupby('stock_code'):
            group = group.sort_values('date').set_index('date')
            # 创建完整的日期范围
            date_range = pd.date_range(start=group.index.min(), end=group.index.max(), freq='D')
            # 重新索引并向前填充
            group_filled = group.reindex(date_range).fillna(method='ffill')
            group_filled['stock_code'] = stock_code
            equity_filled.append(group_filled.reset_index())

        if equity_filled:
            equity_filled_df = pd.concat(equity_filled, ignore_index=True)
            equity_filled_df.rename(columns={'index': 'date'}, inplace=True)
        else:
            logger.error("没有有效的股本数据")
            return

        # 合并价格、股本和股票信息数据
        self.merged_data = pd.merge(
            self.price_data,
            equity_filled_df[['stock_code', 'date', 'capitalization']],
            on=['stock_code', 'date'],
            how='inner'
        )

        # 合并股票信息（上市日期）
        self.stock_info['ipo_date'] = pd.to_datetime(self.stock_info['ipo_date']).dt.tz_localize(None)
        self.merged_data = pd.merge(
            self.merged_data,
            self.stock_info[['stock_code', 'ipo_date']],
            on='stock_code',
            how='inner'
        )

        # 计算上市天数
        self.merged_data['trading_days_since_ipo'] = (
            (self.merged_data['date'] - self.merged_data['ipo_date']).dt.days
        )

        # 合并基本面数据（如果有）
        if self.fundamental_data is not None and len(self.fundamental_data) > 0:
            # 对基本面数据进行向前填充
            fundamental_filled = []
            for stock_code, group in self.fundamental_data.groupby('stock_code'):
                group = group.sort_values('date').set_index('date')
                # 创建完整的日期范围
                date_range = pd.date_range(start=group.index.min(), end=group.index.max(), freq='D')
                # 重新索引并向前填充
                group_filled = group.reindex(date_range).fillna(method='ffill')
                group_filled['stock_code'] = stock_code
                fundamental_filled.append(group_filled.reset_index())

            if fundamental_filled:
                fundamental_filled_df = pd.concat(fundamental_filled, ignore_index=True)
                fundamental_filled_df.rename(columns={'index': 'date'}, inplace=True)

                # 合并基本面数据
                self.merged_data = pd.merge(
                    self.merged_data,
                    fundamental_filled_df[['stock_code', 'date', 'pe_ttm', 'pb', 'ps_ttm', 'dyr', 'to_r', 'mc', 'cmc', 'ecmc']],
                    on=['stock_code', 'date'],
                    how='left'
                )
                logger.info("成功合并基本面数据")

        # 计算市值
        self.merged_data['market_cap'] = self.merged_data['close'] * self.merged_data['capitalization']

        # 移除无效数据
        self.merged_data = self.merged_data.dropna(subset=['market_cap'])
        self.merged_data = self.merged_data[self.merged_data['market_cap'] > 0]

        logger.info(f"合并后数据: {len(self.merged_data)} 条记录")
        logger.info(f"数据时间范围: {self.merged_data['date'].min()} 至 {self.merged_data['date'].max()}")

    def get_bse50_constituents(self, date: datetime) -> set:
        """
        获取指定日期的北证50成分股列表

        包含股票代码映射功能：将430/830开头的代码映射为920开头

        Args:
            date: 查询日期

        Returns:
            成分股代码集合（已映射为920开头）
        """
        if not hasattr(self, 'bse50_constituents_data') or len(self.bse50_constituents_data) == 0:
            logger.warning("没有北证50成分股数据")
            return set()

        date_str = pd.to_datetime(date).strftime('%Y-%m-%d')

        # 找到小于等于指定日期的最近一次成分股变更记录
        valid_changes = self.bse50_constituents_data[
            self.bse50_constituents_data['date'] <= date_str
        ]

        if len(valid_changes) == 0:
            logger.warning(f"日期 {date_str} 没有找到北证50成分股数据")
            return set()

        # 获取最新的成分股记录
        latest_change = valid_changes.iloc[-1]
        constituents_json = latest_change['constituents']

        try:
            import json
            constituents = json.loads(constituents_json)

            # 应用股票代码映射：430/830开头 -> 920开头
            mapped_constituents = set()
            for code in constituents:
                if isinstance(code, str):
                    if code.startswith(('430', '830')):
                        # 将430/830开头的代码映射为920开头
                        mapped_code = '920' + code[3:]
                        mapped_constituents.add(mapped_code)
                    else:
                        # 保持其他代码不变
                        mapped_constituents.add(code)
                else:
                    logger.warning(f"北证50成分股代码格式异常: {code}")

            # 检查映射后的成分股是否在当前股票数据中存在
            available_stocks = set(self.merged_data['stock_code'].unique())
            valid_constituents = mapped_constituents.intersection(available_stocks)

            if len(valid_constituents) == 0:
                logger.warning(f"日期 {date_str}: 映射后的北证50成分股与当前股票数据不匹配")
                logger.warning(f"原始成分股样本: {list(constituents)[:5]}")
                logger.warning(f"映射后成分股样本: {list(mapped_constituents)[:5]}")
                logger.warning(f"当前股票代码样本: {list(available_stocks)[:5]}")
            else:
                logger.info(f"日期 {date_str}: 北证50成分股映射成功 {len(constituents)} -> {len(valid_constituents)}")

            return valid_constituents
        except json.JSONDecodeError as e:
            logger.error(f"解析北证50成分股JSON失败: {e}")
            return set()

    def apply_fundamental_filters(self, data: pd.DataFrame, date: datetime) -> pd.DataFrame:
        """
        应用基本面过滤条件

        Args:
            data: 股票数据DataFrame
            date: 计算日期

        Returns:
            过滤后的DataFrame
        """
        if len(data) == 0:
            return data

        original_count = len(data)
        filtered_data = data.copy()

        # 应用基本面过滤条件
        for indicator, condition in self.filter_conditions.items():
            if indicator in filtered_data.columns:
                min_val = condition.get('min')
                max_val = condition.get('max')

                if min_val is not None:
                    before_count = len(filtered_data)
                    filtered_data = filtered_data[filtered_data[indicator] >= min_val]
                    after_count = len(filtered_data)
                    if before_count > after_count:
                        logger.debug(f"过滤条件 {indicator} >= {min_val}: {before_count} -> {after_count}")

                if max_val is not None:
                    before_count = len(filtered_data)
                    filtered_data = filtered_data[filtered_data[indicator] <= max_val]
                    after_count = len(filtered_data)
                    if before_count > after_count:
                        logger.debug(f"过滤条件 {indicator} <= {max_val}: {before_count} -> {after_count}")

        # 应用上市天数过滤
        if self.min_trading_days is not None and 'trading_days_since_ipo' in filtered_data.columns:
            before_count = len(filtered_data)
            filtered_data = filtered_data[filtered_data['trading_days_since_ipo'] >= self.min_trading_days]
            after_count = len(filtered_data)
            if before_count > after_count:
                logger.debug(f"上市天数过滤 >= {self.min_trading_days}天: {before_count} -> {after_count}")

        # 应用北证50成分股过滤
        if self.filter_bse50_constituents:
            before_count = len(filtered_data)
            bse50_constituents = self.get_bse50_constituents(date)
            if bse50_constituents:
                filtered_data = filtered_data[filtered_data['stock_code'].isin(bse50_constituents)]
                after_count = len(filtered_data)
                if before_count > after_count:
                    logger.debug(f"北证50成分股过滤: {before_count} -> {after_count}")
            else:
                logger.warning(f"日期 {pd.to_datetime(date).strftime('%Y-%m-%d')} 没有北证50成分股数据，跳过此过滤")

        filtered_count = len(filtered_data)
        if filtered_count < original_count:
            date_str = pd.to_datetime(date).strftime('%Y-%m-%d')
            logger.info(f"日期 {date_str}: 总过滤 {original_count} -> {filtered_count} 只股票")

        return filtered_data

    def set_filter_condition(self, indicator: str, min_val: float = None, max_val: float = None):
        """
        设置单个过滤条件

        Args:
            indicator: 指标名称（如 'pe_ttm', 'pb' 等）
            min_val: 最小值（None表示不限制）
            max_val: 最大值（None表示不限制）
        """
        if indicator not in self.filter_conditions:
            self.filter_conditions[indicator] = {'min': None, 'max': None}

        self.filter_conditions[indicator]['min'] = min_val
        self.filter_conditions[indicator]['max'] = max_val

        conditions = []
        if min_val is not None:
            conditions.append(f"min={min_val}")
        if max_val is not None:
            conditions.append(f"max={max_val}")

        logger.info(f"设置过滤条件 {indicator}: {', '.join(conditions) if conditions else '清除'}")

    def set_multiple_filters(self, filters: Dict[str, Dict[str, float]]):
        """
        批量设置过滤条件

        Args:
            filters: 过滤条件字典
                    例如: {'pe_ttm': {'min': 0, 'max': 50}, 'pb': {'min': 0.5, 'max': 5}}
        """
        for indicator, condition in filters.items():
            self.set_filter_condition(indicator, condition.get('min'), condition.get('max'))

        logger.info(f"批量设置 {len(filters)} 个过滤条件")

    def clear_all_filters(self):
        """清除所有过滤条件"""
        for indicator in self.filter_conditions:
            self.filter_conditions[indicator] = {'min': None, 'max': None}
        logger.info("已清除所有过滤条件")

    def set_min_trading_days(self, days: int):
        """
        设置最小上市天数过滤条件

        Args:
            days: 最小上市天数
        """
        self.min_trading_days = days
        logger.info(f"设置最小上市天数: {days}天")

    def set_bse50_filter(self, enable: bool):
        """
        设置是否启用北证50成分股过滤

        Args:
            enable: True表示只选择北证50成分股，False表示不限制
        """
        self.filter_bse50_constituents = enable
        status = "启用" if enable else "禁用"
        logger.info(f"{status}北证50成分股过滤")

    def get_filter_summary(self) -> str:
        """获取当前过滤条件摘要"""
        active_filters = []
        for indicator, condition in self.filter_conditions.items():
            min_val = condition.get('min')
            max_val = condition.get('max')
            if min_val is not None or max_val is not None:
                conditions = []
                if min_val is not None:
                    conditions.append(f">={min_val}")
                if max_val is not None:
                    conditions.append(f"<={max_val}")
                active_filters.append(f"{indicator}({' & '.join(conditions)})")

        # 添加上市天数过滤条件
        if self.min_trading_days is not None:
            active_filters.append(f"上市天数>={self.min_trading_days}天")

        # 添加北证50成分股过滤条件
        if self.filter_bse50_constituents:
            active_filters.append("北证50成分股")

        if active_filters:
            return f"当前过滤条件: {' | '.join(active_filters)}"
        else:
            return "当前无过滤条件"

    def calculate_market_cap_ranking(self, date: datetime) -> pd.DataFrame:
        """
        计算指定日期的市值排名（包含基本面过滤）

        Args:
            date: 计算日期

        Returns:
            按市值排序的股票DataFrame
        """
        # 获取指定日期的数据
        date_data = self.merged_data[self.merged_data['date'] == date].copy()

        if len(date_data) == 0:
            date_str = pd.to_datetime(date).strftime('%Y-%m-%d')
            logger.warning(f"日期 {date_str} 没有数据")
            return pd.DataFrame()

        # 数据质量检查
        valid_data = date_data[date_data['market_cap'] > 0]
        invalid_count = len(date_data) - len(valid_data)

        if invalid_count > 0:
            date_str = pd.to_datetime(date).strftime('%Y-%m-%d')
            logger.warning(f"日期 {date_str} 有 {invalid_count} 只股票市值为0或无效")

        # 应用基本面过滤条件
        filtered_data = self.apply_fundamental_filters(valid_data, date)

        if len(filtered_data) == 0:
            date_str = pd.to_datetime(date).strftime('%Y-%m-%d')
            logger.warning(f"日期 {date_str} 基本面过滤后没有符合条件的股票")
            return pd.DataFrame()

        # 按市值排序
        filtered_data = filtered_data.sort_values('market_cap', ascending=True)
        filtered_data['rank'] = range(1, len(filtered_data) + 1)

        return filtered_data

    def select_stocks(self, date: datetime) -> List[str]:
        """
        选择指定日期的N只最小市值股票

        Args:
            date: 选择日期

        Returns:
            选中的股票代码列表
        """
        ranked_stocks = self.calculate_market_cap_ranking(date)

        if len(ranked_stocks) == 0:
            return []

        # 选择前N只股票
        selected_stocks = ranked_stocks.head(self.N)['stock_code'].tolist()

        date_str = pd.to_datetime(date).strftime('%Y-%m-%d')
        logger.info(f"日期 {date_str} 选中股票: {selected_stocks}")
        return selected_stocks

    def calculate_transaction_cost(self, trade_value: float) -> float:
        """
        计算交易成本（手续费 + 滑点）

        Args:
            trade_value: 交易金额

        Returns:
            交易成本
        """
        commission = trade_value * self.commission_rate
        slippage = trade_value * self.slippage_rate
        return commission + slippage

    def execute_trade(self, date: datetime, stock_code: str, shares: int,
                     price: float, action: str) -> Dict:
        """
        执行交易

        Args:
            date: 交易日期
            stock_code: 股票代码
            shares: 交易股数
            price: 交易价格
            action: 交易类型 ('buy' 或 'sell')

        Returns:
            交易记录
        """
        trade_value = shares * price
        cost = self.calculate_transaction_cost(trade_value)

        if action == 'buy':
            net_cost = trade_value + cost
        else:
            net_cost = trade_value - cost

        trade_record = {
            'date': date,
            'stock_code': stock_code,
            'action': action,
            'shares': shares,
            'price': price,
            'trade_value': trade_value,
            'cost': cost,
            'net_cost': net_cost
        }

        self.trades.append(trade_record)
        return trade_record

    def rebalance_portfolio(self, date: datetime, selected_stocks: List[str]):
        """
        调仓

        Args:
            date: 调仓日期
            selected_stocks: 选中的股票代码列表
        """
        date_str = pd.to_datetime(date).strftime('%Y-%m-%d')
        logger.info(f"开始调仓: {date_str}")

        # 获取调仓日的价格数据
        price_data_date = self.price_data[self.price_data['date'] == date]

        # 卖出不在新组合中的股票
        stocks_to_sell = []
        for stock_code in self.positions:
            if stock_code not in selected_stocks:
                stocks_to_sell.append(stock_code)

        for stock_code in stocks_to_sell:
            position = self.positions[stock_code]
            stock_price_data = price_data_date[price_data_date['stock_code'] == stock_code]

            if len(stock_price_data) > 0:
                sell_price = stock_price_data.iloc[0]['open']  # 按开盘价卖出
                trade_record = self.execute_trade(date, stock_code, position['shares'], sell_price, 'sell')
                logger.info(f"卖出 {stock_code}: {position['shares']}股 @ {sell_price:.2f}, 交易成本: {trade_record['cost']:.2f}")

            del self.positions[stock_code]

        # 计算新组合的等权重资金分配
        if selected_stocks:
            # 计算当前总资产（现金 + 持仓市值）
            total_value = self.initial_capital
            for stock_code, position in self.positions.items():
                stock_price_data = price_data_date[price_data_date['stock_code'] == stock_code]
                if len(stock_price_data) > 0:
                    current_price = stock_price_data.iloc[0]['close']
                    total_value += position['shares'] * current_price

            # 等权重分配
            equal_weight = total_value / len(selected_stocks)

            # 买入新股票
            for stock_code in selected_stocks:
                if stock_code not in self.positions:
                    stock_price_data = price_data_date[price_data_date['stock_code'] == stock_code]

                    if len(stock_price_data) > 0:
                        buy_price = stock_price_data.iloc[0]['open']  # 按开盘价买入
                        shares = int(equal_weight / buy_price / 100) * 100  # 按手买入

                        if shares > 0:
                            trade_record = self.execute_trade(date, stock_code, shares, buy_price, 'buy')
                            self.positions[stock_code] = {
                                'shares': shares,
                                'cost_price': buy_price
                            }
                            logger.info(f"买入 {stock_code}: {shares}股 @ {buy_price:.2f}, 总成本: {trade_record['net_cost']:.2f}")

    def calculate_portfolio_value(self, date: datetime) -> float:
        """
        计算组合总值

        Args:
            date: 计算日期

        Returns:
            组合总值
        """
        total_value = self.initial_capital  # 初始资金

        # 计算持仓市值
        price_data_date = self.price_data[self.price_data['date'] == date]

        for stock_code, position in self.positions.items():
            stock_price_data = price_data_date[price_data_date['stock_code'] == stock_code]
            if len(stock_price_data) > 0:
                current_price = stock_price_data.iloc[0]['close']
                total_value += position['shares'] * current_price

        return total_value

    def run_backtest(self, start_date: str = None, end_date: str = None):
        """
        运行回测

        Args:
            start_date: 开始日期 (YYYY-MM-DD)
            end_date: 结束日期 (YYYY-MM-DD)
        """
        logger.info("开始回测...")
        logger.info(f"策略参数: N={self.N}, M={self.M}, 初始资金={self.initial_capital:,}")

        # 确定回测时间范围
        if start_date is None:
            start_date = self.merged_data['date'].min().tz_localize(None)
        else:
            start_date = pd.to_datetime(start_date).tz_localize(None)

        if end_date is None:
            end_date = self.merged_data['date'].max().tz_localize(None)
        else:
            end_date = pd.to_datetime(end_date).tz_localize(None)

        logger.info(f"回测时间范围: {start_date} 至 {end_date}")

        # 获取所有交易日
        trading_days = sorted(self.merged_data['date'].unique())
        # 转换为datetime类型以确保比较正确
        trading_days = [pd.to_datetime(day) for day in trading_days]
        trading_days = [day for day in trading_days if start_date <= day <= end_date]

        logger.info(f"交易日数量: {len(trading_days)}")

        # 初始化
        self.positions = {}
        self.trades = []
        self.portfolio_value = []
        self.daily_returns = []

        # 调仓计数器
        rebalance_counter = 0

        # 主回测循环
        for i, date in enumerate(trading_days):
            try:
                # 检查是否需要调仓
                if i == 0 or rebalance_counter >= self.M:
                    # 选择股票
                    selected_stocks = self.select_stocks(date)

                    if selected_stocks:
                        # 执行调仓
                        self.rebalance_portfolio(date, selected_stocks)
                        rebalance_counter = 0
                    else:
                        logger.warning(f"日期 {date} 无有效股票选择")
                else:
                    rebalance_counter += 1

                # 计算组合价值
                portfolio_value = self.calculate_portfolio_value(date)
                self.portfolio_value.append({
                    'date': date,
                    'value': portfolio_value
                })

                # 计算日收益率
                if i > 0:
                    prev_value = self.portfolio_value[-2]['value']
                    daily_return = (portfolio_value - prev_value) / prev_value
                    self.daily_returns.append({
                        'date': date,
                        'return': daily_return
                    })

                # 进度输出
                if (i + 1) % 50 == 0:
                    logger.info(f"回测进度: {i + 1}/{len(trading_days)} ({(i + 1)/len(trading_days)*100:.1f}%)")

            except Exception as e:
                date_str = pd.to_datetime(date).strftime('%Y-%m-%d')
                logger.error(f"日期 {date_str} 处理失败: {e}")
                continue

        logger.info("回测完成!")

    def calculate_performance_metrics(self, use_actual_start_date: bool = True) -> Dict:
        """
        计算策略表现指标

        Args:
            use_actual_start_date: 是否使用实际第一次买入日期计算年化收益率
                                  True - 使用第一次买入日期
                                  False - 使用回测开始日期

        Returns:
            表现指标字典
        """
        if not self.portfolio_value:
            return {}

        # 转换为DataFrame
        portfolio_df = pd.DataFrame(self.portfolio_value)
        returns_df = pd.DataFrame(self.daily_returns) if self.daily_returns else pd.DataFrame()

        # 基本指标
        initial_value = self.initial_capital
        final_value = portfolio_df['value'].iloc[-1]
        total_return = (final_value - initial_value) / initial_value

        # 计算时间指标
        backtest_start_date = portfolio_df['date'].iloc[0]
        backtest_end_date = portfolio_df['date'].iloc[-1]

        # 确定用于年化计算的起始日期
        if use_actual_start_date and self.trades:
            # 从交易记录中找到第一次买入日期
            buy_trades = [trade for trade in self.trades if trade['action'] == 'buy']
            if buy_trades:
                first_buy_date = min(trade['date'] for trade in buy_trades)
                investment_start_date = pd.to_datetime(first_buy_date)
                logger.info(f"实际投资开始日期: {investment_start_date.strftime('%Y-%m-%d')} (第一次买入)")
            else:
                logger.warning("没有买入交易记录，使用回测开始日期")
                investment_start_date = backtest_start_date
        else:
            investment_start_date = backtest_start_date
            logger.info(f"使用回测开始日期: {investment_start_date.strftime('%Y-%m-%d')}")

        end_date = portfolio_df['date'].iloc[-1]
        days = (end_date - investment_start_date).days
        years = days / 365.25

        # 年化收益率（基于实际投资期间）
        annual_return = (final_value / initial_value) ** (1/years) - 1 if years > 0 else 0

        # 同时计算基于回测期间的年化收益率（用于对比）
        backtest_days = (backtest_end_date - backtest_start_date).days
        backtest_years = backtest_days / 365.25
        backtest_annual_return = (final_value / initial_value) ** (1/backtest_years) - 1 if backtest_years > 0 else 0

        # 日收益率统计
        if not returns_df.empty:
            daily_returns = returns_df['return']
            volatility = daily_returns.std() * np.sqrt(252)  # 年化波动率
            sharpe_ratio = annual_return / volatility if volatility > 0 else 0

            # 最大回撤
            peak = portfolio_df['value'].expanding().max()
            drawdown = (portfolio_df['value'] - peak) / peak
            max_drawdown = drawdown.min()

            # 最大回撤期间
            max_dd_end = drawdown.idxmin()
            max_dd_start = portfolio_df.loc[:max_dd_end, 'value'].idxmax()

            # 转换为datetime格式
            if hasattr(max_dd_start, 'strftime'):
                max_dd_start = pd.to_datetime(max_dd_start)
            else:
                max_dd_start = pd.to_datetime(portfolio_df.loc[max_dd_start, 'date'])

            if hasattr(max_dd_end, 'strftime'):
                max_dd_end = pd.to_datetime(max_dd_end)
            else:
                max_dd_end = pd.to_datetime(portfolio_df.loc[max_dd_end, 'date'])
        else:
            volatility = 0
            sharpe_ratio = 0
            max_drawdown = 0
            max_dd_start = None
            max_dd_end = None

        # 交易统计
        total_trades = len(self.trades)
        total_cost = sum(trade['cost'] for trade in self.trades)
        cost_ratio = total_cost / initial_value if initial_value > 0 else 0

        # 年度收益率
        portfolio_df['year'] = portfolio_df['date'].dt.year
        yearly_returns = []
        for year in portfolio_df['year'].unique():
            year_data = portfolio_df[portfolio_df['year'] == year]
            if len(year_data) > 1:
                year_start = year_data['value'].iloc[0]
                year_end = year_data['value'].iloc[-1]
                year_return = (year_end - year_start) / year_start
                yearly_returns.append({'year': year, 'return': year_return})

        metrics = {
            'initial_capital': initial_value,
            'final_value': final_value,
            'total_return': total_return,
            'annual_return': annual_return,
            'backtest_annual_return': backtest_annual_return,  # 基于回测期间的年化收益率
            'volatility': volatility,
            'sharpe_ratio': sharpe_ratio,
            'max_drawdown': max_drawdown,
            'max_drawdown_start': max_dd_start,
            'max_drawdown_end': max_dd_end,
            'total_trades': total_trades,
            'total_cost': total_cost,
            'cost_ratio': cost_ratio,
            'trading_days': len(portfolio_df),
            'years': years,
            'backtest_years': backtest_years,  # 基于回测期间的年数
            'backtest_start_date': backtest_start_date,
            'backtest_end_date': backtest_end_date,
            'investment_start_date': investment_start_date,
            'yearly_returns': yearly_returns
        }

        return metrics

    def print_performance_summary(self, use_actual_start_date: bool = True, show_comparison: bool = True):
        """
        打印策略表现摘要

        Args:
            use_actual_start_date: 是否使用实际第一次买入日期计算年化收益率
            show_comparison: 是否显示两种年化收益率的对比
        """
        metrics = self.calculate_performance_metrics(use_actual_start_date)

        if not metrics:
            logger.error("没有有效的回测结果")
            return

        print("\n" + "="*60)
        print("小市值轮动策略回测结果")
        print("="*60)

        print(f"\n基本信息:")
        print(f"  初始资金: ¥{metrics['initial_capital']:,.0f}")
        print(f"  最终资金: ¥{metrics['final_value']:,.0f}")
        print(f"  总收益率: {metrics['total_return']:.2%}")

        # 年化收益率信息
        print(f"\n年化收益率:")
        print(f"  基于实际投资期间: {metrics['annual_return']:.2%}")
        print(f"    投资期间: {metrics['investment_start_date'].strftime('%Y-%m-%d')} 至 {metrics['backtest_end_date'].strftime('%Y-%m-%d')}")
        print(f"    投资时长: {metrics['years']:.1f}年")

        # 显示对比信息
        if show_comparison and (metrics['investment_start_date'] != metrics['backtest_start_date']):
            print(f"  基于回测数据期间: {metrics['backtest_annual_return']:.2%}")
            print(f"    回测期间: {metrics['backtest_start_date'].strftime('%Y-%m-%d')} 至 {metrics['backtest_end_date'].strftime('%Y-%m-%d')}")
            print(f"    回测时长: {metrics['backtest_years']:.1f}年")

            # 计算差异
            diff = metrics['annual_return'] - metrics['backtest_annual_return']
            print(f"  年化收益率差异: {diff:+.2%} (实际投资期间 {'更真实' if diff != 0 else '相同'})")

        print(f"\n交易统计:")
        print(f"  交易日数: {metrics['trading_days']}天")
        print(f"  总交易次数: {len(self.trades)}")

        if self.trades:
            first_trade = min(self.trades, key=lambda x: x['date'])
            last_trade = max(self.trades, key=lambda x: x['date'])
            print(f"  交易期间: {first_trade['date'].strftime('%Y-%m-%d')} 至 {last_trade['date'].strftime('%Y-%m-%d')}")

        print(f"\n风险指标:")
        print(f"  年化波动率: {metrics['volatility']:.2%}")
        print(f"  夏普比率: {metrics['sharpe_ratio']:.2f}")
        print(f"  最大回撤: {metrics['max_drawdown']:.2%}")

        if metrics['max_drawdown_start'] and metrics['max_drawdown_end']:
            try:
                start_date = pd.to_datetime(metrics['max_drawdown_start']).strftime('%Y-%m-%d')
                end_date = pd.to_datetime(metrics['max_drawdown_end']).strftime('%Y-%m-%d')
                print(f"  最大回撤期间: {start_date} 至 {end_date}")
            except Exception as e:
                print(f"  最大回撤期间: 数据格式异常")
                logger.warning(f"最大回撤日期格式化失败: {e}")

        print(f"\n交易统计:")
        print(f"  总交易次数: {metrics['total_trades']}")
        print(f"  总交易成本: ¥{metrics['total_cost']:,.0f}")
        print(f"  成本占比: {metrics['cost_ratio']:.2%}")

        print(f"\n年度收益率:")
        for year_data in metrics['yearly_returns']:
            print(f"  {year_data['year']}: {year_data['return']:.2%}")

        print("\n" + "="*60)

    def plot_results(self):
        """绘制回测结果图表"""
        if not self.portfolio_value:
            logger.error("没有有效的回测数据")
            return

        # 创建图表
        fig, axes = plt.subplots(2, 2, figsize=(15, 10))
        fig.suptitle('小市值轮动策略回测结果', fontsize=16)

        # 转换为DataFrame
        portfolio_df = pd.DataFrame(self.portfolio_value)
        returns_df = pd.DataFrame(self.daily_returns) if self.daily_returns else pd.DataFrame()

        # 1. 组合价值曲线
        axes[0, 0].plot(portfolio_df['date'], portfolio_df['value'], 'b-', linewidth=2)
        axes[0, 0].axhline(y=self.initial_capital, color='r', linestyle='--', alpha=0.7, label='初始资金')
        axes[0, 0].set_title('组合价值变化')
        axes[0, 0].set_ylabel('组合价值 (¥)')
        axes[0, 0].legend()
        axes[0, 0].grid(True, alpha=0.3)

        # 2. 累计收益率
        if len(portfolio_df) > 0:
            cumulative_returns = (portfolio_df['value'] / self.initial_capital - 1) * 100
            axes[0, 1].plot(portfolio_df['date'], cumulative_returns, 'g-', linewidth=2)
            axes[0, 1].axhline(y=0, color='r', linestyle='--', alpha=0.7)
            axes[0, 1].set_title('累计收益率')
            axes[0, 1].set_ylabel('收益率 (%)')
            axes[0, 1].grid(True, alpha=0.3)

        # 3. 日收益率分布
        if not returns_df.empty:
            axes[1, 0].hist(returns_df['return'] * 100, bins=50, alpha=0.7, color='skyblue', edgecolor='black')
            axes[1, 0].axvline(x=returns_df['return'].mean() * 100, color='red', linestyle='--', label=f"均值: {returns_df['return'].mean()*100:.2f}%")
            axes[1, 0].set_title('日收益率分布')
            axes[1, 0].set_xlabel('日收益率 (%)')
            axes[1, 0].set_ylabel('频次')
            axes[1, 0].legend()
            axes[1, 0].grid(True, alpha=0.3)

        # 4. 回撤曲线
        if len(portfolio_df) > 0:
            peak = portfolio_df['value'].expanding().max()
            drawdown = (portfolio_df['value'] - peak) / peak * 100
            axes[1, 1].fill_between(portfolio_df['date'], drawdown, 0, alpha=0.3, color='red')
            axes[1, 1].plot(portfolio_df['date'], drawdown, 'r-', linewidth=1)
            axes[1, 1].set_title('回撤曲线')
            axes[1, 1].set_ylabel('回撤 (%)')
            axes[1, 1].grid(True, alpha=0.3)

        # 调整布局
        plt.tight_layout()

        # 保存图表
        chart_path = 'small_cap_strategy_results.png'
        plt.savefig(chart_path, dpi=300, bbox_inches='tight')
        plt.show()

        logger.info(f"图表已保存: {chart_path}")

    def save_results(self, filename: str = "strategy_results.json"):
        """保存回测结果到文件"""
        results = {
            'strategy_params': {
                'N': self.N,
                'M': self.M,
                'initial_capital': self.initial_capital,
                'commission_rate': self.commission_rate,
                'slippage_rate': self.slippage_rate
            },
            'performance_metrics': self.calculate_performance_metrics(),
            'trades': [
                {
                    'date': pd.to_datetime(trade['date']).strftime('%Y-%m-%d'),
                    'stock_code': trade['stock_code'],
                    'action': trade['action'],
                    'shares': trade['shares'],
                    'price': trade['price'],
                    'trade_value': trade['trade_value'],
                    'cost': trade['cost']
                }
                for trade in self.trades[-100:]  # 只保存最后100条交易记录
            ]
        }

        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(results, f, ensure_ascii=False, indent=2, default=str)

        logger.info(f"回测结果已保存: {filename}")


def main():
    """主函数"""
    try:
        # 检查数据库文件
        db_path = "bse_data.db"
        if not os.path.exists(db_path):
            logger.error(f"数据库文件不存在: {db_path}")
            logger.error("请先运行 bse_data_fetcher.py 获取数据")
            return

        # 创建策略实例
        strategy = SmallCapRotationStrategy(db_path)

        # 可选：自定义参数
        strategy.N = 5  # 选择5只股票
        strategy.M = 10  # 每10个交易日调仓
        strategy.initial_capital = 1000000  # 100万初始资金

        # 运行回测
        # 可指定时间范围，如：strategy.run_backtest("2023-01-01", "2024-12-31")
        strategy.run_backtest()

        # 打印结果（自动使用第一次买入日期作为实际投资开始日期）
        # 这会自动检测交易记录中的第一次买入日期来计算年化收益率
        strategy.print_performance_summary(use_actual_start_date=True, show_comparison=True)

        # 绘制图表
        strategy.plot_results()

        # 保存结果
        strategy.save_results()

    except Exception as e:
        logger.error(f"策略运行失败: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()